# Kestrel: An XMPP-based Many-Task Computing Scheduler
# Author: Lance Stout <lancestout@gmail.com>
#
# Copyright 2009-2010 Clemson University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


"""
Job Queue Module

The job queue module tracks the state of user submitted jobs.
"""

from __future__ import with_statement
import threading
import kestrel.core.datatypes as datatypes
import time

class JobQueue(object):
    """
    The job queue module tracks the state of user submitted jobs.

    Events listened for by this module:
      job/approved
      job/cancelled
      job/task/start
      job/task/end
      job/task/rejected
      schedule/worker

    Events raised by this module:
      task/send
    """

    def __init__(self, kernel):
        """
        Initialize the job queue trackint the state
        of user submitted jobs.
        """
        self.kernel = kernel
        self.events = {"job/approved": self.approved,
                       "job/cancelled": self.cancel_job,
                       "job/task/start": self.task_started,
                       "job/task/end": self.task_finished,
                       "job/task/rejected": self.task_rejected,
                       "job/task/terminated": self.task_terminated,
                       "schedule/worker": self.schedule_worker,
                       "status/requested": self.status}
        self.kernel.register(self)
        self.queue = []
        self.jobs = {}
        self.lock = threading.RLock()


    def approved(self, job):
        """
        A job has been approved, so add it to the queue.
        """
        with self.lock:
            self.jobs[job.job_id] = job
            self.queue.append(job)
            self.kernel.event('schedule/job', job)


    def task_started(self, task):
        """Mark a task as running."""
        with self.lock:
            if task.job_id in self.jobs:
                self.jobs[task.job_id].start_task(task.task_id)
            else:
                # Hack to make sure the worker pool has updated before canceling
                time.sleep(5)
                self.kernel.event('job/task/cancel', {'task_id': task.task_id,
                                                      'job_id': task.job_id})


    def task_finished(self, task):
        """Mark a task as completed."""
        with self.lock:
            if task.job_id in self.jobs:
                job = self.jobs[task.job_id]
                job.complete_task(task.task_id)
                if job.is_completed():
                    self.queue.remove(job)
                    del self.jobs[job.job_id]
                    self.kernel.event('job/finished', job)


    def task_rejected(self, task):
        """
        Return a task to the queue if it has been rejected
        by a worker agent.
        """
        with self.lock:
            if task.job_id in self.jobs:
                job = self.jobs[task.job_id]
                job.requeue_task(task.task_id)


    def task_terminated(self, task):
        """
        Return a task to the queue if the worker agent
        executing it has gone offline.
        """
        with self.lock:
            if task.job_id in self.jobs:
                job = self.jobs[task.job_id]
                job.requeue_task(task.task_id)
                self.kernel.event('schedule/job', job)


    def cancel_job(self, data):
        """Mark a job as canceled."""
        with self.lock:
            job_id = data.get('job_id', None)
            user = data['source']
            user_jobs = []

            # Select jobs to be canceled
            if job_id is None:
                user_jobs = [self.jobs[i] for i in self.jobs \
                                 if self.jobs[i].source == user]
            else:
                user_jobs = [self.jobs[job_id]]

            # Cancel each job
            for job in user_jobs:
                job.cancel()
                if job.is_completed():
                    if job in self.queue:
                        self.queue.remove(job)
                    del self.jobs[job.job_id]
                    self.kernel.event('job/finished', job)

                # There are active tasks running, cancel those
                else:
                    for task_id in job.tasks['running']:
                        self.kernel.event('job/task/cancel', {'job_id': job.job_id,
                                                              'task_id': task_id})

    def schedule_worker(self, worker):
        """Attempt to dispatch a task to the given worker."""
        with self.lock:
            for job_id in self.jobs:
                job = self.jobs[job_id]
                if worker.matches(job.requires):
                    task = job.next_task()
                    if task is None:
                        continue
                    worker.task = ('pending', task)
                    self.kernel.event('task/send', {'to': worker.id,
                                                    'task': task})


    def status(self, data):
        """HACK version to send status"""
        with self.lock:
            def make_stat(job):
                return {'job_id': job.job_id,
                        'running': len(job.tasks['running']),
                        'completed': len(job.tasks['complete']),
                        'queued': job.queue}

            job_id = data.get('job_id', None)

            if job_id is None:
                jobs = [make_stat(self.jobs[i]) for i in self.jobs
                        if self.jobs[i].source == data['source']]

                self.kernel.event('status/send', {'to': data['source'],
                                                  'status': jobs,
                                                  'type':'jobs'})
            else:
                job = self.jobs.get(job_id, None)
                if job is not None and job.source == data['source']:
                    resp = make_stat(job)
                    self.kernel.event('status/send', {'to': data['source'],
                                                      'status': resp,
                                                      'type': 'job'})
                else:
                    resp = 'No such job. It could have finished executing.'
                    self.kernel.event('status/send', {'to': data['source'],
                                                      'status': resp,
                                                      'type':'error'})

